from copy import copy
from dataclasses import dataclass, fields
from itertools import count
from transformers import AutoConfig

@dataclass
class SamplingParams:
    temperature: float = 1.0
    max_tokens: int = 64
    ignore_eos: bool = False

    @property
    def len(self):
        return 1 + self.max_tokens

@dataclass
class Config:
    model: str | None = None
    max_num_batched_tokens: int = 16384 # 单个批处理中允许的最大 token 数量,控制批处理的总大小，防止内存溢出
    max_num_seqs: int = 512 # 单个批处理中允许的最大序列（请求）数量
    max_model_len: int = 4096 # 模型支持的最大上下文长度
    gpu_memory_utilization: float = 0.9
    tensor_parallel_size: int = 1
    enforce_eager: bool = False # 控制是否跳过PyTorch的图编译优化 True: 调试模式，便于调试但性能较低 False: 生产模式，启用优化但可能遇到编译错误
    hf_config: AutoConfig | None = None # HuggingFace 模型配置对象 获取模型的架构细节（层数、隐藏维度等）
    eos: int = -1 #  -1 表示使用模型配置中的默认EOS
    kvcache_block_size: int = 256
    num_kvcache_blocks: int = -1

    def __post_init__(self):
        assert os.path.isdir(self.model)
        assert self.kvcache_block_size % 256 == 0
        assert 1 <= self.tensor_parallel_size <= 8
        self.hf_config = AutoConfig.from_pretrained(self.model)
        self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
        assert self.max_num_batched_tokens >= self.max_model_len

class LLMEngine:
    def __init__(self, **kwargs):
        config_fields = {field.name for field in fields(Config)}
        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
        config = Config(**config_kwargs)




if __name__ == '__main__':
    print(129 // 16)
    print((129 // 16) + 1)
    print((130 + 16 - 1) // 16)
    print((128 + 16 - 1) // 16)
    sampling_params = SamplingParams()
    print(f"sampling_params.len: {sampling_params.len}")
    prompts = list()
    prompts.append(1)
    prompts.append(2)
    prompts.append(3)
    print(f"prompts[:-1]: {prompts[:-1]}]")
    if not isinstance(sampling_params, list):
        sampling_params = [sampling_params] * len(prompts)
    for prompt, sp in zip(prompts, sampling_params):
        print(prompt, sp)

    counter = count()
    print(next(counter))
    print(next(counter))

    str1 = "str1"
    print(copy(str1))

    dict1 = {"name": "linzimng", "age": 1}
    dict2 = {"name": "linhai", "age": 2}
    dict2.update(dict1)
    print(dict2)
    dict3 = {**dict1, **dict2}
    print(dict3)
    LLMEngine(enforce_eager=True, tensor_parallel_size=1)